import torch
import torch.nn as nn
from torchvision import transforms, datasets
from torchvision.models import resnet50
from torch.utils.data import DataLoader
from PIL import Image


# 要实现输入一张图片，并在给定的数据集中找到最相似的图片，可以使用预训练的DINOv2模型和图像相似度度量方法。以下是一个示例代码：
# 请确保将"数据集路径"替换为实际的图像数据集路径，将"预训练模型权重路径"替换为预训练的DINOv2模型权重的路径，
# 将"查询图片路径"替换为要查询的图片的路径。这段代码中，我们首先加载查询图片并将其转换为张量，然后将其输入到DINOv2模型中以提取特征向量。
# 接下来，我们计算数据集中每张图片与查询图片的相似度，使用余弦相似度作为相似度度量方法。最后，我们找到具有最高相似度的图片索引，并打印出其路径作为结果。

# 定义DINOv2模型结构
class DINOv2(nn.Module):
    def __init__(self, num_classes):
        super(DINOv2, self).__init__()
        self.encoder = resnet50(pretrained=True)
        self.head = nn.Linear(2048, num_classes)

    def forward(self, x):
        x = self.encoder(x)
        x = x.view(x.size(0), -1)
        x = self.head(x)
        return x


# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义数据预处理的转换
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载数据集
dataset = datasets.ImageFolder("./images", transform=transform)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 创建DINOv2模型实例
model = DINOv2(num_classes=len(dataset.classes))
model.to(device)

# 加载预训练的DINOv2模型权重
state_dict = torch.load("./model/dinov2_vitl14_pretrain.pth")
model.load_state_dict(state_dict, strict=False)
model.eval()

# 加载要查询的图片
query_image_path = "./images/penguin.jpg"
print("输入图片路径：", query_image_path)
query_image = Image.open(query_image_path).convert("RGB")
query_image = transform(query_image).unsqueeze(0).to(device)

# 提取查询图片的特征向量
with torch.no_grad():
    query_features = model.encoder(query_image)

# 计算数据集中每张图片与查询图片的相似度
similarities = []
with torch.no_grad():
    for images, _ in dataloader:
        images = images.to(device)
        features = model.encoder(images)
        similarities.extend(torch.cosine_similarity(query_features, features, dim=1).cpu())

# 存储每个相似度分数及其对应的图片路径
similarity_paths = []

# 按照相似度分数从大到小排序并打印结果
for index, similarity_score in enumerate(similarities):
    image_path = dataset.imgs[index][0]
    similarity_paths.append((similarity_score, image_path))

# 按照相似度分数从高到低排序并打印每个相似度分数及其对应的图片路径
sorted_similarity_paths = sorted(similarity_paths, key=lambda x: x[0], reverse=True)
for similarity_score, image_path in sorted_similarity_paths:
    print(f'相似度分数 {similarity_score:.2f}: {image_path}')